import os

import torch
import torch.nn as nn
from torchvision.models import resnet50
import torch_pruning as tp

from models.lwCETModel import lwCET
from models.models import ecgTransForm
from myutils.modelUtils import get_configs

# 1. 加载预训练模型
exp_log_dir = '../experiments_logs/Brn/ecg-trans-268_2025_10_09_21_20'
test_model_name = os.path.join(exp_log_dir, "checkpoint_best.pt")
dataset_configs, hparams_class1 = get_configs()
hparams = hparams_class1.train_params
# model = lwCET(dataset_configs, hparams).eval()
original_model = ecgTransForm(dataset_configs, hparams).eval()
model = ecgTransForm(dataset_configs, hparams).eval()

input1 = torch.randn(128, 1, 268)  # 第一个输入，例如ECG信号
# input2 = torch.randn(128, 52)      # 第二个输入，例如特征向量
# example_inputs = (input1, input2)
example_inputs = input1

# 2. 构建依赖图（自动分析模型结构）
DG = tp.DependencyGraph()
DG.build_dependency(model, example_inputs=example_inputs)

# 3. 获取所有可剪枝的分组（自动识别）
ignored_layers = []
for module in model.modules():
    if isinstance(module, nn.Linear) and module.out_features == 1000:  # 分类头
        ignored_layers.append(module)

pruning_groups = DG.get_all_groups(ignored_layers=ignored_layers)

# 4. 定义剪枝策略并执行
pruning_ratio = 0.5  # 最终目标剪枝比例
num_iterations = 5    # 迭代次数（例如5次）
fine_tune_epochs = 1  # 每次剪枝后的微调轮数

# 创建优化器和损失函数
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

for group in pruning_groups:
    if len(group) == 1:
        continue
    # 计算要剪枝的通道数（至少保留1个通道）
    prune_count = max(1, int(len(group) * pruning_ratio))
    prune_count = min(prune_count, len(group) - 1)  # 确保不超过最大可剪枝数
    # 生成索引列表
    idxs = list(range(prune_count))
    group.prune(idxs=idxs)

# 检查剪枝后模型是否可运行
try:
    output = model(example_inputs)
    print("✅ 自动剪枝成功！模型前向传播正常。")
except Exception as e:
    print(f"❌ 剪枝后模型前向传播失败: {e}")
    # 可以考虑回滚操作或使用其他剪枝策略


# 计算剪枝前后的参数量
def count_parameters(models):
    return sum(p.numel() for p in models.parameters() if p.requires_grad)


print(f"剪枝前参数量: {count_parameters(original_model)}")
print(f"剪枝后参数量: {count_parameters(model)}")
